/*
 * Dump DWARF var/func address.
 *
 * Copyright (c) 2024 laokz <zhangkai@iscas.ac.cn>
 * Klpmake is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of
 * the Mulan PSL v2. You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES
 * OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
 * TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */
#define _XOPEN_SOURCE 500
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <gelf.h>
#include <libdwarf-0/dwarf.h>
#include <libdwarf-0/libdwarf.h>

/* only dwarf call, hnn ..., need predefined 'err' variable */
#define E(statement, fmt, ...) do {                                 \
    int ret_ = statement;                                           \
    if (ret_ == DW_DLV_NO_ENTRY) {                                  \
        /* unexpected, DWARF info might corrupted */                \
        fprintf(stderr, "ERROR: %s:%d no entry: "fmt"\n", __func__, __LINE__,\
                                            ## __VA_ARGS__);        \
        return ret_;                                        \
    } else if (ret_ == DW_DLV_ERROR) {                              \
        fprintf(stderr, "ERROR: %s:%d %s\n", __func__, __LINE__,    \
                                            dwarf_errmsg(err));     \
        return ret_;                                        \
    }                                                               \
} while (0)

#define KIND_CU 0
#define KIND_VAR 1
#define KIND_FUNC 2

static int get_var_addr(Dwarf_Attribute attr, Dwarf_Addr *addr)
{
    Dwarf_Unsigned lcount = 0;
    Dwarf_Loc_Head_c loclist_head = 0;
    Dwarf_Error err;
    unsigned dw_lkind;

    E(dwarf_get_loclist_c(attr, &loclist_head, &lcount, &err), "");
    /*
     * Only care "Single location descriptions",
     * not debug usage "Location lists".
     */
    E(dwarf_get_loclist_head_kind(loclist_head, &dw_lkind, &err), "");
    if (dw_lkind != DW_LKIND_expression)
        return DW_DLV_NO_ENTRY;

    /* don't care so many args, but that would fail if miss them */
    Dwarf_Small loclist_lkind = 0;
    Dwarf_Small lle_value = 0;
    Dwarf_Unsigned rawval1 = 0;
    Dwarf_Unsigned rawval2 = 0;
    Dwarf_Bool debug_addr_unavailable = 0;
    Dwarf_Addr lopc = 0;
    Dwarf_Addr hipc = 0;
    Dwarf_Unsigned loclist_expr_op_count = 0;
    Dwarf_Locdesc_c locdesc_entry = 0;
    Dwarf_Unsigned expression_offset = 0;
    Dwarf_Unsigned locdesc_offset = 0;
    E(dwarf_get_locdesc_entry_d(loclist_head, 0, &lle_value, &rawval1, &rawval2,
                &debug_addr_unavailable, &lopc, &hipc, &loclist_expr_op_count,
                &locdesc_entry, &loclist_lkind, &expression_offset,
                &locdesc_offset, &err), "");
    /* not care stack or whatever debug usage */
    if (loclist_expr_op_count !=1 )
        return DW_DLV_NO_ENTRY;

    Dwarf_Small op = 0;
    Dwarf_Unsigned opd1 = 0;
    Dwarf_Unsigned opd2 = 0;
    Dwarf_Unsigned opd3 = 0;
    Dwarf_Unsigned offsetforbranch = 0;
    /* location has only one entry */
    E(dwarf_get_location_op_value_c(locdesc_entry, 0, &op, addr, &opd2, &opd3,
                        &offsetforbranch, &err), "");

    dwarf_dealloc_loc_head_c(loclist_head);
    if (op == DW_OP_addr)
        return DW_DLV_OK;
    return DW_DLV_NO_ENTRY;
}

static int get_addr(Dwarf_Debug dbg, Dwarf_Die die, Dwarf_Addr *addr, int kind)
{
    Dwarf_Signed atcount;
    Dwarf_Attribute *atlist;
    Dwarf_Half attrnum = 0;
    Dwarf_Error err;
    int ret = DW_DLV_NO_ENTRY;

    E(dwarf_attrlist(die, &atlist, &atcount, &err), "");
    for (int i = 0; i < atcount; ++i) {
        E(dwarf_whatattr(atlist[i], &attrnum, &err), "");
        if ((kind == KIND_VAR) && (attrnum == DW_AT_location)) {
            ret = get_var_addr(atlist[i], addr);
            break;
        } else if ((kind == KIND_FUNC) && (attrnum == DW_AT_low_pc)) {
            E(dwarf_formaddr(atlist[i], addr, &err), "");
            ret = DW_DLV_OK;
            break;
        }
        dwarf_dealloc_attribute(atlist[i]);
    }
    dwarf_dealloc(dbg, atlist, DW_DLA_LIST);
    return ret;
}

static inline int has_die_name(Dwarf_Die die, char **diename)
{
    return dwarf_diename(die, diename, NULL) == DW_DLV_OK;
}

static void dump_name_addr(char *diename, Dwarf_Addr addr, int kind)
{
    if (kind == KIND_CU)
        printf("-%s\n", diename);
    else if (kind == KIND_VAR)
        printf("+%s %lx\n", diename, addr);
    else
        printf("*%s %lx\n", diename, addr);
}

static int dump_cu_ranges(Dwarf_Debug dbg, Dwarf_Die cu)
{
    Dwarf_Signed atcount;
    Dwarf_Attribute *atlist;
    Dwarf_Half attrnum = 0;
    Dwarf_Error err;

    E(dwarf_attrlist(cu, &atlist, &atcount, &err), "");
    for (int i = 0; i < atcount; ++i) {
        E(dwarf_whatattr(atlist[i], &attrnum, &err), "");
        if (attrnum == DW_AT_ranges) {  /* non-continuous address ranges */
            Dwarf_Ranges *r;
            Dwarf_Off off;
            Dwarf_Signed ranges_count;

            /* DW_AT_ranges point to an offset to .debug_ranges section */
            E(dwarf_global_formref(atlist[i], &off, &err), "");
            E(dwarf_get_ranges_b(dbg, off, cu, NULL, &r, &ranges_count, NULL, &err), "");
            /* here not check FORM and RANGE kind? */
            for(int k = 0; k < ranges_count; k++) {
                printf("%lx %lx\n", r[k].dwr_addr1, r[k].dwr_addr2);
            }
            dwarf_dealloc_ranges(dbg, r, ranges_count);
            break;
        } else if (attrnum == DW_AT_low_pc) {   /* single continuous address range */
            Dwarf_Half form = 0;
            enum Dwarf_Form_Class formclass = DW_FORM_CLASS_UNKNOWN;
            Dwarf_Addr lo, hi;
            E(dwarf_lowpc(cu, &lo, &err), "");
            /* highpc might be an address, or an offset to lowpc */
            E(dwarf_highpc_b(cu, &hi, &form, &formclass, &err), "");
            if ((form != DW_FORM_addr) && !dwarf_addr_form_is_indexed(form)) {
                hi += lo;
            }
            printf("%lx %lx\n", lo, hi);
            break;
        }
        dwarf_dealloc_attribute(atlist[i]);
    }
    dwarf_dealloc(dbg, atlist, DW_DLA_LIST);
    return DW_DLV_OK;
}

static int wall_through_recursive(Dwarf_Debug dbg, Dwarf_Die parent)
{
    Dwarf_Die kid, d;
    Dwarf_Half tag;
    Dwarf_Addr addr;
    Dwarf_Error err;
    char *diename;
    int ret;

    if ((ret = dwarf_child(parent, &d, &err)) == DW_DLV_NO_ENTRY) {
        return DW_DLV_OK;
    } else if (ret == DW_DLV_ERROR) {
        fprintf(stderr, "ERROR: %s:%d %s\n", __func__, __LINE__, dwarf_errmsg(err));
        return ret;
    }
    kid = NULL;
    do {
        if (kid)
            dwarf_dealloc_die(kid);
        E(dwarf_tag(d, &tag, &err), "");
        if ((tag == DW_TAG_variable) && has_die_name(d, &diename) &&
            /*
             * When module is statically compiled into the kernel,
             * module_exit() has no effect, then its __exitcall_...
             * static variable address is 0.
             * So does __UNIQUE_ID_...?
             * Ignore them.
             */
            ((ret = get_addr(dbg, d, &addr, KIND_VAR)) == DW_DLV_OK) && addr) {
            printf("\t");
            dump_name_addr(diename, addr, KIND_VAR);
        } else {
            ret = wall_through_recursive(dbg, d);
        }
        kid = d;
        if (ret == DW_DLV_ERROR)
            break;
        ret = dwarf_siblingof_b(dbg, kid, 1, &d, &err);
    } while(ret == DW_DLV_OK);

    dwarf_dealloc_die(kid);
    if (ret == DW_DLV_NO_ENTRY)
        ret = DW_DLV_OK;
    return ret;
}

static int wall_through_cu(Dwarf_Debug dbg, Dwarf_Die cu, int has_ranges)
{
    Dwarf_Die kid, d;
    Dwarf_Error err;
    Dwarf_Addr addr;
    char *diename;
    int ret, kind;

    if (!has_die_name(cu, &diename))
        return ret;
    dump_name_addr(diename, 0, KIND_CU);
    if (has_ranges && ((ret = dump_cu_ranges(dbg, cu)) != DW_DLV_OK))
        return ret;

    if ((ret = dwarf_child(cu, &kid, &err)) == DW_DLV_ERROR) {
        fprintf(stderr, "ERROR: %s:%d %s\n", __func__, __LINE__, dwarf_errmsg(err));
        return ret;
    } else if (ret == DW_DLV_NO_ENTRY) {    /* head.S has no child */
        return DW_DLV_OK;
    }
    while((ret = dwarf_siblingof_b(dbg, kid, 1, &d, &err)) == DW_DLV_OK) {
        Dwarf_Half tag;
        char *diename;
        E(dwarf_tag(d, &tag, &err), "");
        if ((tag == DW_TAG_variable) || (tag == DW_TAG_subprogram)) {
            kind = tag == DW_TAG_variable ? KIND_VAR : KIND_FUNC;
            if (has_die_name(d, &diename)) {
                ret = get_addr(dbg, d, &addr, kind);
                if ((ret == DW_DLV_OK) && addr) {
                    dump_name_addr(diename, addr, kind);
                    if (kind == KIND_FUNC)
                        ret = wall_through_recursive(dbg, d);
                }
            }
            if (ret == DW_DLV_ERROR) {
                dwarf_dealloc_die(d);
                break;
            }
        }

        dwarf_dealloc_die(kid);
        kid = d;
    }

    if (ret == DW_DLV_ERROR)
        fprintf(stderr, "ERROR: %s:%d %s\n", __func__, __LINE__, dwarf_errmsg(err));
    else
        ret = DW_DLV_OK;
    dwarf_dealloc_die(kid);
    return ret;
}

static int wall_through(Dwarf_Debug dbg, int has_ranges)
{
    Dwarf_Die cu_die;
    Dwarf_Error err;
    int ret;

    do {
        ret = dwarf_next_cu_header_d(dbg, 1, NULL, NULL, NULL, NULL, NULL,
                    NULL, NULL, NULL, NULL, NULL, &err);
        if (ret == DW_DLV_NO_ENTRY) {
            break;
        } else if (ret == DW_DLV_ERROR) {
            fprintf(stderr, "ERROR: %s:%d %s\n", __func__, __LINE__, dwarf_errmsg(err));
            return ret;
        }

        E(dwarf_siblingof_b(dbg, NULL, 1, &cu_die, &err), "");
        ret = wall_through_cu(dbg, cu_die, has_ranges);
        dwarf_dealloc_die(cu_die);
    } while (ret == DW_DLV_OK);

    return ret;
}

#define EE(statement, err_value) do {   \
    if ((statement) == (err_value)) {   \
        fprintf(stderr, "ERROR: %s:%d %s", __func__, __LINE__, elf_errmsg(0));\
        ret = -1;                       \
        goto out;                       \
    }                                   \
} while (0)
static int has_debug_ranges_sec(char *mod)
{
    Elf *relf = NULL;
    int rfd, ret;
    size_t sections_nr, shstrtab_index, i;
    char *name;
    Elf_Scn *scn = NULL;
    GElf_Shdr sh;

    EE(rfd = open(mod, O_RDONLY), -1);
    EE(elf_version(EV_CURRENT), EV_NONE);
    EE(relf = elf_begin(rfd, ELF_C_READ_MMAP_PRIVATE, NULL), NULL);
    EE(elf_getshdrnum(relf, &sections_nr), -1);
    EE(elf_getshdrstrndx(relf, &shstrtab_index), -1);
    for (i = 1; i < sections_nr; i++) {
        EE(scn = elf_nextscn(relf, scn), NULL);
        EE(gelf_getshdr(scn, &sh), NULL);
        EE(name = elf_strptr(relf, shstrtab_index, sh.sh_name), NULL);
        if (!strcmp(name, ".debug_ranges"))
            break;
    }
    ret = i != sections_nr;

out:
    if (relf)
        elf_end(relf);
    close(rfd);
    return ret;
}

int main(int argc, char *argv[])
{
    Dwarf_Debug dbg;
    Dwarf_Error err;
    int flag;

    if (argc != 2) {
        fprintf(stderr, "dump variable and function address in DWARF\n");
        fprintf(stderr, "Usage: $0 vmlinux-or-module-with-debuginfo\n");
        exit(1);
    }

    flag = has_debug_ranges_sec(argv[1]);
    if (flag == -1)
        flag = 0;
    E(dwarf_init_path(argv[1], NULL, 0, DW_GROUPNUMBER_ANY, NULL, NULL,
                                &dbg, &err), "no DWARF info");
    (void)wall_through(dbg, flag);
    dwarf_finish(dbg);
}
